import zarr
import numpy as np
import os
import sys
from pathlib import Path
from loadJson import _load_finger_length
def convert_14_16(data_14,embodiment_name=None):
    if embodiment_name is not None:
        orig_action_dim=_load_finger_length()[embodiment_name]["action_space"]
        if orig_action_dim==16:
            assert data_14.shape[-1] == 16, f"Expected last dim 16, got {data_14.shape}"
            return data_14
    assert data_14.shape[-1] == 14, f"Expected last dim 14, got {data_14.shape}"
    shape_16 = list(data_14.shape)
    print(f"data_14.shape: {data_14.shape}")
    shape_16[-1] = 16
    data_16 = np.zeros(shape_16, dtype=data_14.dtype)
    data_16[..., :6] = data_14[..., :6]
    data_16[..., 6] = 0
    data_16[..., 7:14] = data_14[..., 6:13]
    data_16[..., 14] = 0
    data_16[..., 15] = data_14[..., 13]
    return data_16
def convert_16_14(data_16,embodiment_name):
    orig_action_dim=_load_finger_length()[embodiment_name]["action_space"]
    assert data_16.shape[-1] == 16, f"Expected last dim 16, got {data_16.shape}"
    if orig_action_dim==16:
        return data_16
    if not np.allclose(data_16[..., 6], 0,atol=0.05):
        print(f"Position 6 should be 0, got {data_16[..., 6]} {data_16}")
    if not np.allclose(data_16[..., 14], 0,atol=0.05):
        print(f"Position 14 should be 0, got {data_16[..., 14]} {data_16}")
    shape_14 = list(data_16.shape)
    shape_14[-1] = 14
    data_14 = np.zeros(shape_14, dtype=data_16.dtype)
    data_14[..., :6] = data_16[..., :6]
    data_14[..., 6:13] = data_16[..., 7:14]
    data_14[..., 13] = data_16[..., 15]
    assert data_14.shape[-1] == 14, f"Expected last dim 14, got {data_14.shape}"
    return data_14
def process_zarr(zarr_path):
    zarr_path = Path(zarr_path)
    assert os.path.exists(zarr_path), f"ZARR path does not exist: {zarr_path}"
    zarr_group = zarr.open(zarr_path, mode='r+')
    for key in ["action", "state"]:
        js_key_path = os.path.join(zarr_path, 'data', f'js_{key}')
        assert os.path.exists(js_key_path), f"js_{key} path does not exist: {js_key_path}"
        js_key_source = Path(js_key_path).readlink()
        source_path = zarr_path/"data"/js_key_source
        print(f"source_path: {source_path}")
        source_group = zarr.open(source_path, mode='r+')
        js_key_data = source_group
        assert js_key_data.shape[-1] == 14, f"{key} last dim should be 14, got {js_key_data.shape}"
        js_key_16 = convert_14_16(js_key_data)
        chunks_16 = list(js_key_data.chunks)
        chunks_16[-1] = 16
        print(f"orig {list(np.round(js_key_data[10], 3))}")
        print(f"new {list(np.round(js_key_16[10], 3))}")
        input(f"js_key_data.shape: {js_key_data.shape}, js_key_16.shape: {js_key_16.shape} chunks_16: {chunks_16}. continue?")
        os.rename(str(source_path), str(source_path)+"_14")
        zarr_group[source_path.parent.name].create_dataset(f'{key}', data=js_key_16, chunks=chunks_16, compressor=js_key_data.compressor, dtype=js_key_data.dtype, fill_value=js_key_data.fill_value, order=js_key_data.order)
if __name__ == "__main__":
    assert len(sys.argv) == 2, f"Usage: {sys.argv[0]} <zarr_path>"
    zarr_path = sys.argv[1]
    process_zarr(zarr_path) 